import paho.mqtt.client as mqtt
import subprocess
import sys
import time
import platform
import os
import json
import logging
import requests
from threading import Timer
from hashlib import md5
from datetime import datetime

g_config = {}
with open("config.json", "r") as f:
    g_config = json.load(f)

MQTT_OFFLINE_TIMEOUT = 3600

COMMAND_TOPIC = f"shell/{g_config['device_name']}/cmd"
RESPONSE_TOPIC = f"shell/{g_config['device_name']}/rsp"

def update_mqtt_broker_addr():
    new_addr = ""
    try:
        response = requests.get("http://free.frp.vip:10272/files/addr.json")
        if response.status_code == 200:
            new_addr = response.json().get("mqtt_addr", "")
    except requests.RequestException as e:
        print(f"获取MQTT Broker地址失败: {e}")
        return
    
    if (new_addr and new_addr != g_config['mqtt_addr']):
        g_config['mqtt_addr'] = new_addr
        logging.info(f"update mqtt broker addr to: {new_addr}")
        with open("config.json", "w") as f:
            json.dump(g_config, f, indent=4)

class MQTTSalve:
    def __init__(self):
        # 初始化日志
        self.setup_logging()
        
        self.client = mqtt.Client()
        self.client.on_connect = self.on_connect
        self.client.on_disconnect = self.on_disconnect
        self.client.on_message = self.on_message
        self.mqtt_online_time = 0
        self.mqtt_offline_time = int(time.time())
        self.reboot_timer = None
        self.reboot_timer = Timer(MQTT_OFFLINE_TIMEOUT, self.offline_timeout_cb)
        self.reboot_timer.start()
        self.is_windows = platform.system() == "Windows"
        # 定义危险命令列表（排除restart和reboot，因为它们有特殊处理）
        self.dangerous_commands = [
            'rm -rf', 'format', 'del /f', 'rmdir /s', 
            'shutdown', 'kill',
            'dd if=/dev/zero', 'mkfs', 'fdisk'
        ]
        # 当前工作目录
        self.current_directory = os.getcwd()
        
    def setup_logging(self):
        """设置日志记录"""
        # 创建logs目录（如果不存在）
        log_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "logs")
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)
            
        # 获取今天的日期作为日志文件名
        today = datetime.now().strftime("%Y-%m-%d")
        log_file = os.path.join(log_dir, f"slave_{today}.log")
        
        # 配置日志格式
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file, encoding='utf-8'),
                logging.StreamHandler(sys.stdout)  # 同时输出到控制台
            ]
        )
        
        self.logger = logging.getLogger("MQTTSalve")
        
    def on_connect(self, client, userdata, flags, rc):
        if rc == 0:
            # 订阅命令主题
            self.client.subscribe(COMMAND_TOPIC, qos=1)
            self.logger.info(f"mqtt connect success, subscribe topic: {COMMAND_TOPIC}")
            self.mqtt_online_time = int(time.time())
            if self.reboot_timer:
                self.reboot_timer.cancel()
        else:
            self.logger.error(f"mqtt connect failed: {rc}")
            self.mqtt_offline_time = int(time.time())

    def on_disconnect(self, client, userdata, rc):
        self.logger.error(f"mqtt disconnected: {rc}")
        self.mqtt_offline_time = int(time.time())
        self.reboot_timer = Timer(MQTT_OFFLINE_TIMEOUT, self.offline_timeout_cb)
        self.reboot_timer.start()
    
    def offline_timeout_cb(self):
        self.logger.info("mqtt offline timeout")
        update_mqtt_broker_addr()
        self.restart_client()

    def generate_s_data(self, pos=0):
        rand = (int(time.time()) + pos * 3) // 3
        rand_code = g_config["rand_code"]
        s_data = f"{rand_code}{rand}"
        s = md5(s_data.encode('utf-8')).hexdigest()
        return s[:8]

    def verify_s_data(self, s):
        s1 = self.generate_s_data(0) # 当前s
        s2 = self.generate_s_data(-1) # 前一个s
        return s1 == s or s2 == s

    def on_message(self, client, userdata, msg):
        if msg.topic == COMMAND_TOPIC:
            try:
                # 解析JSON格式的命令
                json_command = json.loads(msg.payload.decode('utf-8'))
                s = json_command.get("s", "")
                if not self.verify_s_data(s):
                    self.logger.warning("s verify failed")
                    self.send_response("s verify failed")
                    return
                command = json_command.get("data", "")
            except json.JSONDecodeError:
                # 如果不是JSON格式，拒绝执行
                self.logger.warning("data format error")
                self.send_response("data format error")
                return
            
            self.logger.info(f"[收到命令]: {command}")
            
            # 检查是否为重启命令
            if command.strip().lower() == 'restart' or command.strip().lower() == 'reboot':
                self.logger.info("检测到重启命令，正在重启程序...")
                result = self.restart_client()
                self.send_response(result)
                return
            
            # 检查命令安全性
            if self.is_command_safe(command):
                result = self.execute_command(command)
            else:
                error_msg = "错误: 拒绝执行危险命令"
                self.logger.warning(f"拒绝执行危险命令: {command}")
                result = error_msg
                
            self.send_response(result)
    
    def is_command_safe(self, command):
        """检查命令是否安全"""
        # 转换为小写进行检查
        cmd_lower = command.lower()
        
        # 检查是否包含危险命令
        for dangerous in self.dangerous_commands:
            if dangerous in cmd_lower:
                return False
                
        # 检查是否尝试访问敏感目录
        sensitive_paths = ['/etc/', '/root/', '/boot/', '/sys/']
        if not self.is_windows:  # 仅在Unix系统上检查
            for path in sensitive_paths:
                if path in command:
                    return False
                    
        return True
    
    def execute_command(self, command):
        try:
            # 设置超时时间（秒）
            timeout = 30
            
            # 处理 cd 命令
            if command.strip().startswith('cd '):
                return self.change_directory(command)
            
            # 根据操作系统选择shell，并设置工作目录
            if self.is_windows:
                # Windows系统使用cmd，使用二进制模式读取输出以避免编码问题
                process = subprocess.Popen(
                    command,
                    shell=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    cwd=self.current_directory  # 设置工作目录
                )
            else:
                # Unix/Linux/Mac系统使用bash
                process = subprocess.Popen(
                    command,
                    shell=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    executable='/bin/bash',
                    cwd=self.current_directory  # 设置工作目录
                )
            
            # 等待进程完成或超时
            try:
                stdout, stderr = process.communicate(timeout=timeout)
            except subprocess.TimeoutExpired:
                process.kill()
                stdout, stderr = process.communicate()
                timeout_msg = f"错误: 命令执行超时 ({timeout}秒)"
                self.logger.warning(timeout_msg)
                return timeout_msg
            
            # 手动解码输出
            def decode_output(output_bytes):
                if not output_bytes:
                    return ""
                
                # 尝试多种编码解码
                encodings = ['gbk', 'utf-8', 'gb2312']
                for encoding in encodings:
                    try:
                        return output_bytes.decode(encoding)
                    except (UnicodeDecodeError, LookupError):
                        continue
                
                # 如果所有编码都失败，使用错误处理方式
                try:
                    return output_bytes.decode('utf-8', errors='ignore')
                except:
                    return "[无法解码的输出]"
            
            # 解码标准输出和错误输出
            decoded_stdout = decode_output(stdout)
            decoded_stderr = decode_output(stderr)
            
            if process.returncode == 0:
                return decoded_stdout
            else:
                return f"错误 (返回码 {process.returncode}): {decoded_stderr}"
        except Exception as e:
            return f"执行命令时出错: {str(e)}"
    
    def change_directory(self, command):
        """处理 cd 命令并更新当前目录"""
        try:
            # 提取目标目录
            target_dir = command.strip()[3:].strip()  # 去掉 "cd " 前缀
            
            if not target_dir:
                # cd 命令没有参数，返回当前目录
                return self.current_directory
            
            # 处理特殊路径
            if target_dir == '~':
                target_dir = os.path.expanduser('~')
            elif target_dir == '-':
                # 在简单实现中，我们不跟踪上一个目录
                return "错误: 不支持 '-' 参数"
            else:
                # 处理相对路径和绝对路径
                if not os.path.isabs(target_dir):
                    target_dir = os.path.join(self.current_directory, target_dir)
            
            # 解析路径并检查是否存在
            target_dir = os.path.abspath(target_dir)
            
            if os.path.exists(target_dir) and os.path.isdir(target_dir):
                self.current_directory = target_dir
                return f"已切换到目录: {self.current_directory}"
            else:
                return f"错误: 目录不存在: {target_dir}"
        except Exception as e:
            return f"切换目录时出错: {str(e)}"
    
    def restart_client(self):
        """重启当前程序"""
        try:
            self.logger.info("正在重启程序...")
            self.send_response("正在重启程序...")
            
            # 断开MQTT连接
            try:
                self.logger.info("断开MQTT连接")
                self.client.disconnect()
            except Exception as e:
                self.logger.warning(f"断开MQTT连接时出错: {str(e)}")
                pass  # 忽略断开连接时的任何错误
            
            # 重启当前程序
            import sys
            import os
            os.execlp("python", "python", *sys.argv)
        except Exception as e:
            error_msg = f"重启程序时出错: {str(e)}"
            print(error_msg)
            self.logger.error(error_msg)
            return error_msg
    
    def send_response(self, response):
        # 确保响应不是空的
        if not response:
            response = "[无输出]"
            
        # 添加当前目录信息到响应
        if response == "[无输出]":
            response = f"[当前目录: {self.current_directory}]\n{response}"
        else:
            response = f"[当前目录: {self.current_directory}]\n{response}"
            
        # 如果响应过长，截取前1000个字符
        if len(response) > 1000:
            response = response[:1000] + "...[响应过长已截断]"
            
        # 将响应格式化为JSON
        json_response = json.dumps({
            "s": "",
            "data": response
        }, ensure_ascii=False)
            
        result = self.client.publish(RESPONSE_TOPIC, json_response, qos=1)
        if result.rc != mqtt.MQTT_ERR_SUCCESS:
            error_msg = f"发送响应失败: {result.rc}"
            self.logger.error(error_msg)
    
    def connect(self):
        try:
            self.logger.info(f"mqtt connect: {g_config['mqtt_addr']}")
            broker_addr = g_config['mqtt_addr'].split(":")[0]
            broker_port = int(g_config['mqtt_addr'].split(":")[1])
            self.client.connect(broker_addr, broker_port, 60)
            self.client.loop_forever()
        except Exception as e:
            self.logger.error(f"mqtt connect err: {e}")
            
    def disconnect(self):
        self.logger.info("断开MQTT连接")
        self.client.disconnect()

if __name__ == "__main__":
    try:
        slave = MQTTSalve()
        try:
            slave.connect()
        except KeyboardInterrupt:
            slave.logger.info("收到键盘中断，退出程序")
            slave.disconnect()
        except Exception as e:
            error_msg = f"发生错误: {e}"
            slave.logger.error(error_msg)
            slave.disconnect()
    except Exception as e:
        logging.error(f"MQTT Slave初始化失败: {e}")
